import datetime
import json

def get_quarter_boundaries(date_obj: datetime.date):
    """
    根据指定日期判断所属季度（此处季度定义为：
    - 月份在 1~4：从当年 01 到 04
    - 月份在 5~8：从当年 05 到 08
    - 月份在 9~12：从当年 09 到 12
    返回该日期所在季度的起始年月字符串，和结束年月字符串
    """
    year = date_obj.year
    month = date_obj.month
    if month in [1, 2, 3, 4]:
        return f"{year}-01", f"{year}-04"
    elif month in [5, 6, 7, 8]:
        return f"{year}-05", f"{year}-08"
    elif month in [9, 10, 11, 12]:
        return f"{year}-09", f"{year}-12"
    
def get_standard_quarter(date_obj: datetime.date) -> int:
    """
    根据给定日期返回标准日历季度（1-3月为Q1，4-6月为Q2，7-9月为Q3，10-12月为Q4）
    """
    month = date_obj.month
    if month <= 3:
        return 1
    elif month <= 6:
        return 2
    elif month <= 9:
        return 3
    else:
        return 4

def transform_json(input_json: dict) -> dict:
    """
    根据dims字段判断需要如何处理startDate和endDate的值：
    - 如果 dims 中包含 "MonthAnaly": 保留年月格式 "YYYY-MM"
    - 如果 dims 中包含 "QuarterAnaly": 根据日期所属的季度（自定义为1-4月、5-8月、9-12月）计算季度起始和结束年月
      如果起始和结束日期不在同一个季度，则返回包含整个跨度的季度起止，
      即 startDate 取起始日期所在季度的起始值，endDate 取结束日期所在季度的结束值
    - 如果 dims 中包含 "YearAnaly": 只保留年份 "YYYY"
    - 否则，返回原输入不作修改。
    """
    dims = input_json.get("dims", "")
    
    # 判断字段中是否包含指定的模式
    if "MonthAnaly" in dims:
        try:
            start_date = datetime.datetime.strptime(input_json.get("startDate", ""), "%Y-%m-%d").date()
            end_date = datetime.datetime.strptime(input_json.get("endDate", ""), "%Y-%m-%d").date()
        except ValueError:
            # 日期格式不正确时直接返回原数据
            return input_json
        # 仅保留年月
        input_json["startDate"] = start_date.strftime("%Y-%m")
        input_json["endDate"] = end_date.strftime("%Y-%m")
    
    elif "QuarterAnaly" in dims:
        try:
            start_date = datetime.datetime.strptime(input_json.get("startDate", ""), "%Y-%m-%d").date()
            end_date = datetime.datetime.strptime(input_json.get("endDate", ""), "%Y-%m-%d").date()
        except ValueError:
            return input_json
        # 根据标准季度规则计算
        q_start = get_standard_quarter(start_date)
        q_end = get_standard_quarter(end_date)
        input_json["startDate"] = f"{start_date.year}-Q{q_start}"
        input_json["endDate"] = f"{end_date.year}-Q{q_end}"

    elif "YearAnaly" in dims:
        try:
            start_date = datetime.datetime.strptime(input_json.get("startDate", ""), "%Y-%m-%d").date()
            end_date = datetime.datetime.strptime(input_json.get("endDate", ""), "%Y-%m-%d").date()
        except ValueError:
            return input_json
        # 保留年份
        input_json["startDate"] = start_date.strftime("%Y")
        input_json["endDate"] = end_date.strftime("%Y")
    
    # 如果dims中不包含以上三种情况，则不做任何修改，直接返回输入JSON
    return input_json

# 示例测试
if __name__ == "__main__":
    # 示例入参
    json_str = '''
    {
      "startDate": "2024-02-01",
      "endDate": "2024-12-31",
      "kpis": "S0010033 as 门急诊药品总收入_核算日期, S0010025 as 全院药品总收入_核算日期",
      "dims": "AccountDate->MonthAnaly as 核算日期",
      "filter": ""
    }
    '''
    data = json.loads(json_str)
    transformed = transform_json(data)
    print("转换后的结果:")
    print(json.dumps(transformed, indent=2, ensure_ascii=False))
